#!/usr/bin/env python3
# black.py
# Generate code from Mali kernel headers

from pprint import pprint

def parse_structs(f):
    level = 0
    current_struct = ""
    mode = ""
    special = False
    command = None
    fields = {}
    out = {}

    for line in f:
        if line.isspace():
            continue

        if level == 0:
            if line.startswith("struct"):
                current_struct = line.split(" ")[1]
                mode = "IN/OUT"
                level += 1

                fields = {"name": current_struct, "IN": [], "IN/OUT": [], "OUT": []}
            elif line.strip() == "// special":
                special = True
            else:
                print("Unknown level 0")
                print(line)
        elif level == 1:
            if line[0] == "}":
                level = 0
                fields["special"] = special
                special = False
                out[current_struct] = fields
            elif line.strip().startswith("//"):
                parts = line.strip().split(" ")
                cmd = parts[1]

                if cmd == "dump" or cmd == "trace" or cmd == "special":
                    command = parts
                else:
                    mode = cmd
            else:
                parts = line.strip()[0:-1].split(" ")

                if parts[-1] == 'header':
                    continue

                arraySize = False

                if parts[-1][-1] == ']':
                    array_sides = parts[-1].split("[")

                    parts[-1] = array_sides[0]
                    arraySize = array_sides[1][0:-1]

                if parts[-1] == 'padding':
                    continue

                if parts[-1][0] == '*':
                    parts[-2] += "*"
                    parts[-1] = parts[-1][1:]

                name = parts[-1];
                type_s = " ".join(parts[0:-1])

                fields[mode].append([type_s, name, arraySize, command])

                if command:
                    command = None
        else:
            print("Unknown level " + str(level))

    return out

formats = {
    "u64": "0x%LX",
    "u32": "0x%X",
    "u16": "0x%hX",
    "u8": "0x%hhX",
    "s64": "%Ld",
    "s32": "%d",
    "s16": "%hd",
    "s8": "%hhd",
    "char": "'%c'",
    "int": "%d",
    "union kbase_pointer": "%p", # TODO: Don't break on 64-bit
    "base_mem_handle": "%p"
}

def printk_field(field, ctx_name, ctx_val, op):
    name = ctx_name + "." + field[1]

    if field[2]:
        if field[0] == 'char':
            return ([name + " = %s"], [field[1] + "_null"], [["//",
                "null", field[1], field[2], ctx_val + op]])
        elif field[2].isdigit():
            accA = []
            accB = []
            accC = []

            bare = field[1]
            l = int(field[2])
            field[2] = None

            for i in range(0, l):
                field[1] = bare + "[" + str(i) + "]"

                (a, b, c) = printk_field(field, ctx_name, ctx_val, op)
                accA += a
                accB += b
                accC += c

            return (accA, accB, accC)
        else:
            return ([name + " = undecoded array"], [], [])
    elif field[0].startswith("struct"):
        inner = field[0][len("struct "):]

        if structs.__contains__(inner):
            accA = []
            accB = []
            accC = []

            for f in structs[inner]["IN/OUT"]:
                (a, b, c) = printk_field(f, name, ctx_val + op + field[1], ".")
                accA += a
                accB += b
                accC += c

            return (accA, accB, accC)
        else:
            return ([name + " = undecoded struct"], None, [])
    else:
        form = [name + " = " + formats[field[0]]]

        orig = ctx_val + op + field[1] 
        value = orig

        if field[0] == "union kbase_pointer":
            value += ".value"

        command = []

        if field[3]:
            command = [field[3] + [value]]

        return (form, [value], command)

def make_printk(arrow, call, fields, ref, io):
    fs_parts = []
    field_v = [call]
    commands = []

    for field in fields:
        (form, val, command) = printk_field(field, "", ref, "->")

        fs_parts += form
        
        if val:
            field_v += val

        if command:
            commands += command

    format_str = '"' + arrow + " = { " + ", ".join(fs_parts) + " };" + '"'

    for cmd in commands:
        if cmd[1] == 'null':
            o = cmd[3]

            if not cmd[3].isdigit():
                o = cmd[4] + o

            print("char *" + cmd[2] + "_null = " + cmd[4] + cmd[2] + ";")
            print(cmd[2] + "_null[" + o + "] = 0;")

    print("printk(" + ", ".join([format_str] + field_v) + ");")

    for cmd in commands:
        print("{")

        if cmd[1] == 'dump':
            print("size_t d_sz = 0;")
            print("uint8_t *dump = kbase_fetch_" + cmd[2] + "(kctx, " +
                    cmd[4] + ", &d_sz, " + cmd[3] + ");")
            print('formatted_hex_dump("' + cmd[4] + '", dump, d_sz);')
            print("kfree(dump);")
        elif cmd[1] == 'trace':
            count = ref + "->" + cmd[4]
            print("int t;")
            print("struct " + cmd[3] + " *it;")
            print("size_t sz = sizeof(*it) * " + count + ";")
            print("it = kbase_fetch_" + cmd[2] + "(kctx, " + cmd[5] + ", NULL, sz);")
            print("for(t = 0; t < " + count + "; ++t) {")
            make_printk(cmd[5] + "[%d]", "t", structs[cmd[3]]["IN/OUT"],
                    "(&it[t])", structs[cmd[3]])
            print("}")
            print("kfree(it);")

        print("}")

    if io["special"]:
        print("kbase_trace_special_" + io["name"] + "(kctx, " + ref + ");")

special_names = {
        "uku_version_check_args": "UKP_FUNC_ID_CHECK_VERSION",
        "kbase_uk_sync_now": "KBASE_FUNC_SYNC",
        "kbase_uk_profiling_controls": "KBASE_FUNC_GET_PROFILING_CONTROLS",
        "kbase_uk_get_ddk_version": "KBASE_FUNC_GET_VERSION",
        "kbase_uk_gpuprops": "KBASE_FUNC_GPU_PROPS_REG_DUMP",
        "kbase_uk_context_id": "KBASE_FUNC_GET_CONTEXT_ID",
        "kbase_uk_ext_buff_kds_data": "KBASE_FUNC_EXT_BUFFER_LOCK"
}

def ioctl_id(call):
    if special_names.__contains__(call):
        return special_names[call]
    else:
        return ("kbase_func_" + call[len("kbase_uk_"):]).upper()

def make_section(arrow, call, io, mode):
    arr = io[mode] + io["IN/OUT"]

    if len(arr):
        make_printk(arrow + " %s *v", '"' + call + '"',  arr, "v", io)

def make_case(call):
    io = ioctls[call]

    print("case " + ioctl_id(call) + ": {")

    if len(io["IN"] + io["OUT"] + io["IN/OUT"]):
        print("struct " + call + " *v = args;");

    print("if(in) {")
    make_section("->", call, io, "IN")
    print("} else {")
    make_section("<-", call, io, "OUT")
    print("}")

    print("break;")
    print("}")

ioctls = {}

with open("mali_ioctl.h") as f:
    ioctls = parse_structs(f)

structs = {}

with open("mali_structs.h") as f:
    structs = parse_structs(f)

print("void kbase_trace_call(struct kbase_context *kctx, void * const args, u32 id, u32 args_size, bool in) {")

print("switch(id) {")

for call in sorted(ioctls.keys()):
    make_case(call)

print('default: { /* printk("Bad ioctl %d", id); */ }')

print("}")
print("}")
